import chromadb
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
client = chromadb.PersistentClient(path="./chromadb_data")


def save_text_to_db(text, collection_name, metadata={}):
    collection = client.get_or_create_collection(collection_name)
    # 使用文本内容的哈希值生成一个唯一的ID
    text_id = str(abs(hash(text)))
    metadata["text_id"] = text_id
    # 生成向量并转成列表的形式
    text_embedding = model.encode(text).tolist()
    # 向集合添加文本 元数据 ID 和向量
    collection.add(
        documents=[text],
        metadatas=[metadata],
        ids=[text_id],
        embeddings=[text_embedding],
    )
    return text_id


""" text_id = save_text_to_db("这是一个测试的文本")
print(f"text_id:{text_id}")
 """
